#!/usr/bin/env python
# -*- coding: utf-8 -*-

"""
最优化计算方法 - 第一章演示
梯度下降法求解无约束优化问题
以二次函数f(x) = x1^2 + 2*x2^2为例
"""

import numpy as np
import matplotlib.pyplot as plt
from matplotlib import cm
from mpl_toolkits.mplot3d import Axes3D

# 定义目标函数
def f(x):
    """
    目标函数 f(x) = x1^2 + 2*x2^2
    """
    return x[0]**2 + 2*x[1]**2

# 定义梯度函数
def grad_f(x):
    """
    目标函数的梯度
    """
    return np.array([2*x[0], 4*x[1]])

# 梯度下降法实现
def gradient_descent(f, grad_f, x0, alpha=0.1, max_iter=100, tol=1e-6):
    """
    梯度下降法求解无约束优化问题
    
    参数:
        f: 目标函数
        grad_f: 梯度函数
        x0: 初始点
        alpha: 步长（学习率）
        max_iter: 最大迭代次数
        tol: 收敛容差
        
    返回:
        x: 最优点
        f_values: 迭代过程中的函数值
        x_history: 迭代过程中的点
    """
    x = x0.copy()
    x_history = [x.copy()]
    f_values = [f(x)]
    
    # 打印初始信息
    print("迭代开始:")
    print(f"迭代次数\t\tf(x)\t\t||∇f(x)||")
    grad = grad_f(x)
    grad_norm = np.linalg.norm(grad)
    print(f"0\t\t{f(x):.6f}\t{grad_norm:.6f}")
    
    # 迭代优化
    for i in range(max_iter):
        # 计算梯度
        grad = grad_f(x)
        grad_norm = np.linalg.norm(grad)
        
        # 如果梯度范数小于阈值，认为收敛
        if grad_norm < tol:
            print(f"\n梯度下降法收敛!")
            break
            
        # 更新位置
        x = x - alpha * grad
        
        # 记录历史
        x_history.append(x.copy())
        f_values.append(f(x))
        
        # 打印当前迭代信息
        print(f"{i+1}\t\t{f(x):.6f}\t{grad_norm:.6f}")
    else:
        print(f"\n达到最大迭代次数，未收敛!")
    
    # 打印最终结果
    print(f"最优解: x1 = {x[0]:.6f}, x2 = {x[1]:.6f}")
    print(f"最优值: f(x) = {f(x):.6f}")
    
    return x, f_values, x_history

# 可视化函数
def visualize_results(f, x_history, f_values):
    """
    可视化优化结果
    
    参数:
        f: 目标函数
        x_history: 迭代过程中的点
        f_values: 迭代过程中的函数值
    """
    # 将列表转为numpy数组，便于处理
    x_history = np.array(x_history)
    
    # 绘制等高线和优化路径
    plt.figure(figsize=(10, 8))
    
    # 创建网格
    x1 = np.linspace(-1, 1, 100)
    x2 = np.linspace(-2, 2, 100)
    X1, X2 = np.meshgrid(x1, x2)
    Z = np.zeros_like(X1)
    
    # 计算每个网格点的函数值
    for i in range(X1.shape[0]):
        for j in range(X1.shape[1]):
            Z[i, j] = f(np.array([X1[i, j], X2[i, j]]))
    
    # 绘制等高线
    plt.contour(X1, X2, Z, 20, cmap='viridis')
    plt.colorbar(label='f(x)')
    
    # 绘制优化路径
    plt.plot(x_history[:, 0], x_history[:, 1], 'r.-', markersize=8, linewidth=1, label='优化路径')
    plt.plot(x_history[0, 0], x_history[0, 1], 'go', markersize=10, label='初始点')
    plt.plot(x_history[-1, 0], x_history[-1, 1], 'bs', markersize=10, label='最优点')
    
    plt.title('梯度下降法优化过程')
    plt.xlabel('$x_1$')
    plt.ylabel('$x_2$')
    plt.legend()
    plt.grid(True)
    
    # 绘制目标函数值随迭代次数的变化
    plt.figure(figsize=(10, 6))
    plt.plot(f_values, 'b.-', linewidth=2)
    plt.title('目标函数值随迭代次数的变化')
    plt.xlabel('迭代次数')
    plt.ylabel('f(x)')
    plt.grid(True)
    
    # 绘制3D曲面
    fig = plt.figure(figsize=(12, 10))
    ax = fig.add_subplot(111, projection='3d')
    surf = ax.plot_surface(X1, X2, Z, cmap=cm.coolwarm, alpha=0.8, linewidth=0, antialiased=True)
    
    # 添加优化路径到3D图
    z_path = [f(point) for point in x_history]
    ax.plot(x_history[:, 0], x_history[:, 1], z_path, 'r.-', markersize=8, linewidth=2, label='优化路径')
    ax.plot([x_history[0, 0]], [x_history[0, 1]], [z_path[0]], 'go', markersize=10, label='初始点')
    ax.plot([x_history[-1, 0]], [x_history[-1, 1]], [z_path[-1]], 'bs', markersize=10, label='最优点')
    
    fig.colorbar(surf, shrink=0.5, aspect=5, label='f(x)')
    ax.set_title('目标函数曲面和优化路径')
    ax.set_xlabel('$x_1$')
    ax.set_ylabel('$x_2$')
    ax.set_zlabel('f(x)')
    ax.legend()
    
    plt.show()

# 主程序
if __name__ == "__main__":
    # 初始点
    x0 = np.array([1.0, 2.0])
    
    # 运行梯度下降法
    x_opt, f_values, x_history = gradient_descent(f, grad_f, x0, alpha=0.1, max_iter=100, tol=1e-6)
    
    # 可视化结果
    visualize_results(f, x_history, f_values) 